{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "# Finding large collision-free configuration-space regions with IRIS\n",
    "\n",
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    BsplineTrajectory,\n",
    "    DiagramBuilder,\n",
    "    InverseKinematics,\n",
    "    IrisInConfigurationSpace,\n",
    "    IrisOptions,\n",
    "    MathematicalProgram,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    Parser,\n",
    "    PositionConstraint,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    Role,\n",
    "    Solve,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import PublishPositionTrajectory\n",
    "from manipulation.scenarios import AddIiwa, AddPlanarIiwa, AddShape, AddWsg\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reaching into the shelves\n",
    "\n",
    "Note that I'm using the original collision geometry (not replacing the hand geometry with the spheres, like we did to help kinematic trajectory optimization)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def AnimateIris(root_context, plant, visualizer, region, speed):\n",
    "    \"\"\"\n",
    "    A simple hit-and-run-style idea for visualizing the IRIS regions:\n",
    "    1. Start at the center. Pick a random direction and run to the boundary.\n",
    "    2. Pick a new random direction; project it onto the current boundary, and run along it. Repeat\n",
    "    \"\"\"\n",
    "\n",
    "    plant_context = plant.GetMyContextFromRoot(root_context)\n",
    "    visualizer_context = visualizer.GetMyContextFromRoot(root_context)\n",
    "\n",
    "    q = region.ChebyshevCenter()\n",
    "    plant.SetPositions(plant_context, q)\n",
    "    visualizer.ForcedPublish(visualizer_context)\n",
    "    active_face = None\n",
    "\n",
    "    print(\"Press the 'Stop Animation' button in Meshcat to continue.\")\n",
    "    meshcat.AddButton(\"Stop Animation\", \"Escape\")\n",
    "\n",
    "    rng = np.random.default_rng()\n",
    "    nq = plant.num_positions()\n",
    "    prog = MathematicalProgram()\n",
    "    qvar = prog.NewContinuousVariables(nq, \"q\")\n",
    "    prog.AddLinearConstraint(\n",
    "        region.A(), 0 * region.b() - np.inf, region.b(), qvar\n",
    "    )\n",
    "    cost = prog.AddLinearCost(np.ones((nq, 1)), qvar)\n",
    "\n",
    "    while meshcat.GetButtonClicks(\"Stop Animation\") < 1:\n",
    "        direction = rng.standard_normal(nq)\n",
    "        cost.evaluator().UpdateCoefficients(direction)\n",
    "\n",
    "        result = Solve(prog)\n",
    "        assert result.is_success()\n",
    "\n",
    "        q_next = result.GetSolution(qvar)\n",
    "\n",
    "        # Animate between q and q_next (at speed):\n",
    "        # TODO: normalize step size to speed... e.g. something like\n",
    "        # 20 * np.linalg.norm(q_next - q) / speed)\n",
    "        for t in np.append(np.arange(0, 1, 0.05), 1):\n",
    "            qs = t * q_next + (1 - t) * q\n",
    "            plant.SetPositions(plant_context, qs)\n",
    "            visualizer.ForcedPublish(visualizer_context)\n",
    "            if running_as_notebook:\n",
    "                time.sleep(0.05)\n",
    "\n",
    "        q = q_next\n",
    "\n",
    "        if not running_as_notebook:\n",
    "            break\n",
    "\n",
    "    meshcat.DeleteButton(\"Stop Animation\")\n",
    "\n",
    "\n",
    "def animate_iris_region():\n",
    "    meshcat.Delete()\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddIiwa(plant)\n",
    "    wsg = AddWsg(plant, iiwa, welded=True, sphere=False)\n",
    "    p_TopShelf = [0.95, 0, 0.65]\n",
    "\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\"package://manipulation/shelves.sdf\")[0]\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"shelves_body\", bin),\n",
    "        RigidTransform([0.88, 0, 0.4]),\n",
    "    )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(role=Role.kIllustration),\n",
    "    )\n",
    "    collision_visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        meshcat,\n",
    "        MeshcatVisualizerParams(\n",
    "            prefix=\"collision\", role=Role.kProximity, visible_by_default=False\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    num_q = plant.num_positions()\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "\n",
    "    # First seed should just be the home position.\n",
    "    options = IrisOptions()\n",
    "    options.num_collision_infeasible_samples = 2 if running_as_notebook else 1\n",
    "    options.random_seed = 1235\n",
    "    options.require_sample_point_is_contained = True\n",
    "    region = IrisInConfigurationSpace(plant, plant_context, options)\n",
    "\n",
    "    # Add a seed for reaching into the top shelf.\n",
    "    ik = InverseKinematics(plant, plant_context)\n",
    "    collision_constraint = ik.AddMinimumDistanceLowerBoundConstraint(\n",
    "        0.001, 0.01\n",
    "    )\n",
    "    grasp_constraint = ik.AddPositionConstraint(\n",
    "        gripper_frame, [0, 0.1, 0], plant.world_frame(), p_TopShelf, p_TopShelf\n",
    "    )\n",
    "\n",
    "    q = ik.q()\n",
    "    prog = ik.get_mutable_prog()\n",
    "    prog.SetInitialGuess(q, q0)\n",
    "    result = Solve(ik.prog())\n",
    "    if not result.is_success():\n",
    "        print(\"IK failed\")\n",
    "    plant.SetPositions(plant_context, result.GetSolution(q))\n",
    "    diagram.ForcedPublish(context)\n",
    "    print(region.PointInSet(result.GetSolution(q)))\n",
    "\n",
    "    options.configuration_obstacles = [region]\n",
    "    shelf_region = IrisInConfigurationSpace(plant, plant_context, options)\n",
    "    print(shelf_region.PointInSet(result.GetSolution(q)))\n",
    "    AnimateIris(context, plant, visualizer, shelf_region, speed=0.1)\n",
    "\n",
    "\n",
    "animate_iris_region()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
